package com.hdu.transmitableThreadLocal;

import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

public class TransmitableThreadLocalHolder {

    private final static Map<Thread, Map<TransmitableThreadLocal<?>, Object>> HOLDER
            = new ConcurrentHashMap<>();

    public static void register(TransmitableThreadLocal<?> threadLocal) {
        HOLDER.computeIfAbsent(Thread.currentThread(), k -> new ConcurrentHashMap<>())
                .put(threadLocal, threadLocal.get());
    }

    public static Map<TransmitableThreadLocal<?>, Object> get() {
        return HOLDER.get(Thread.currentThread());
    }

    public static void unregister(TransmitableThreadLocal<?> transmitableThreadLocal) {
        if (HOLDER.containsKey(Thread.currentThread())) {
            HOLDER.get(Thread.currentThread()).remove(transmitableThreadLocal);
            if (HOLDER.get(Thread.currentThread()).isEmpty()) {
                HOLDER.remove(Thread.currentThread());
            }
        }
    }
}
